//
// Created by lidongwei on 2023/7/12.
//
#include "iostream"
#include "util/pretreatment.h"
#include "util/infer.h"
#include "util/post_processing.h"

int main(){
    int input_channels=4;//4 or 3
    int classes=2;
    int batch=1;
    int output_size=batch*classes*320*480;
    cv::Size feature_map_size=cv::Size(480,320);
    //////////////////预处理/////////////////////
    static std::string img_file_name="../../data/val/imgs/097.jpg";//需要注意的是opencv的引用是在util目录下的文件中，所以这里的图片路径是要相对于它的
    cv::Size resize=cv::Size(480,320);
    std::vector<unsigned char> model=trt_infer::Infer::load_file("../engine/best_checkpoint_epoch162_train0.0023_val0.9939.engine");//读取模型
    auto engine=trt_infer::Infer::create_engine(model,batch);///将模型反序列化

    float object_threshold=0.9;
    uint8_t labels_color[]={0,255};//类别数组
    static std::string img_save_dir="../util";

    for(int i = 0;i<10;i++){
        cv::Mat img=cv::imread(img_file_name);
        img=handle::pretreatment::Affine_transformation(img,resize,0,2);
        auto start=clock();
        std::vector<float> input_data=handle::pretreatment::Normalization(img,255,input_channels);

        //////////////////读取模型以及推理//////////////
        std::vector<float> output_data=trt_infer::Infer::infer(engine,input_data,batch,input_channels,resize.width,resize.height,output_size);

        //////////////////输出结果后处理//////////////
        std::vector<uint8_t> result=handle::post_processing::segment_processing(output_data,classes,feature_map_size,object_threshold,labels_color);//解析出类别
        auto end=clock();
        double diff_time=(double)(end-start)/CLOCKS_PER_SEC*1000;
        printf("整体耗时：%.1f ms\n",diff_time);
        auto result_img=handle::post_processing::drow(result,img,resize);
        cv::imwrite(img_save_dir + "/test.jpg", result_img);
    }


//    try {
//        cv::imwrite(img_save_dir + "/test.jpg", result_img);
//        std::cout << "图片已成功保存到：" << img_save_dir << std::endl;
//    } catch (const cv::Exception& e) {
//        std::cout << "保存图像时出现异常：" << e.what() << std::endl;
//    }

    return 0;
}